from astropy.time import Time
from astropy import units as u
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import os
import imageio.v2 as imageio
from joblib import Parallel, delayed

# 确保输出目录存在
os.makedirs("frames", exist_ok=True)

from pysagin.spacebase import Constellation, Satellite
from pysagin.groundbase import GroundStation
from pysagin.airbase import UavBase, Trajectory
from pysagin.device import RadioFixedAccess, RadioRandomAccess, Radio

# region 初始化全局参数
print("INFO: 初始化全局参数")
start_time = Time("2025-03-27 00:00:00")
# endregion

# region 创建 Walker 星座
print("INFO: 创建 Walker 星座")
constellation = Constellation(
    name="S",
    type="Walker",
    T=66,  # 总卫星数量
    P=6,   # 轨道面数量
    F=0.5, # 相位因子
    a=6371+550,  # 半长轴（单位：km）
    e=0,  # 偏心率
    inc=53,  # 倾角（单位：度）
    init_time=start_time
)
constellation.update_all(start_time)
# endregion

# region 创建地面站
print("INFO: 创建地面站")
station_1 = GroundStation(
    name="G1",
    init_time=start_time,
    init_gps=[126.662892, 45.742140, 0.146]
)
station_2 = GroundStation(
    name="G2",
    init_time=start_time,
    init_gps=[122.116342, 37.510357, 0.007]
)
station_3 = GroundStation(
    name="G3",
    init_time=start_time,
    init_gps=[114.057865, 22.543096, 0.004]
)
stations_gps = [station_1.current_gps, station_2.current_gps, station_3.current_gps]
stations_lon_list = [g[0] for g in stations_gps]
stations_lat_list = [g[1] for g in stations_gps]
# endregion

# region 创建无人机
print("INFO: 创建无人机")
uav_1_trajectory = Trajectory(name="U1Tr",avg_speed=(100*u.km/u.h).to(u.m/u.s).value)
uav_1_trajectory.add_straight(start=station_1.init_gps, end=station_2.init_gps)
uav_1_trajectory.add_figure8(center=station_2.init_gps, size=2000, loops=10)
uav_1_trajectory.add_straight(start=station_2.init_gps, end=station_1.init_gps)
uav_1 = UavBase(
    name="U1",
    init_time=start_time,
    init_gps=station_1.init_gps
)
uav_1.set_trajectory(trajectory=uav_1_trajectory, 
                     departure_time=start_time + 60 * u.min)

uav_2_trajectory = Trajectory(name="U2Tr",avg_speed=(150*u.km/u.h).to(u.m/u.s).value)
uav_2_trajectory.add_straight(start=station_1.init_gps, end=station_3.init_gps)
uav_2_trajectory.add_circle(center=station_3.init_gps, radius=1000, loops=5)
uav_2_trajectory.add_straight(start=station_3.init_gps, end=station_2.init_gps)
uav_2 = UavBase(
    name="U2",
    init_time=start_time,
    init_gps=station_1.init_gps
)
uav_2.set_trajectory(trajectory=uav_2_trajectory, 
                     departure_time=start_time + 120 * u.min)
# endregion

# region 创建无线电台
print("INFO: 创建无线电台")
constellation.create_inter_satellite_links(Device_Radio=RadioFixedAccess, range=10000)

constellation_radio_random = []
for sat in constellation.satellites:
    sat_radio_random = RadioRandomAccess(name=f"{sat.name}D1", platform=sat, range=1000,
                                          max_access_num=-1)
    sat.devices.append(sat_radio_random)
    constellation_radio_random.append(sat_radio_random)

sation_radio_1 = RadioRandomAccess(name=f"{station_1.name}D0", platform=station_1, 
                                   range=1000, max_access_num=-1)
station_1.devices.append(sation_radio_1)
sation_radio_2 = RadioRandomAccess(name=f"{station_2.name}D0", platform=station_2, 
                                   range=1000, max_access_num=-1)
station_2.devices.append(sation_radio_2)
sation_radio_3 = RadioRandomAccess(name=f"{station_3.name}D0", platform=station_3, 
                                   range=1000, max_access_num=-1)
station_3.devices.append(sation_radio_3)

uav_radio_1 = RadioRandomAccess(name=f"{uav_1.name}D0", platform=uav_1, 
                                range=1000, max_access_num=-1)
uav_1.devices.append(uav_radio_1)
uav_radio_2 = RadioRandomAccess(name=f"{uav_2.name}D0", platform=uav_2, 
                                range=1000, max_access_num=-1)
uav_2.devices.append(uav_radio_2)

for sat_radio_random in constellation_radio_random:
    sat_radio_random.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3]+
                                           [uav_radio_1, uav_radio_2])
sation_radio_1.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_2.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_3.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
uav_radio_1.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_2.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)

ALL_DEVICES:list[Radio] = []
for sat in constellation.satellites:
    ALL_DEVICES.extend(sat.devices)
ALL_DEVICES.extend([sation_radio_1, sation_radio_2, sation_radio_3, uav_radio_1, uav_radio_2])
# endregion

# region 生成动画帧（并行化）
print("INFO: 配置动画参数")
duration = 24 * u.hour  # 总时长
interval = 10 * u.min  # 时间间隔
time_steps = np.arange(0, duration.to(u.s).value, interval.to(u.s).value) * u.s
time_series = start_time + time_steps
total_frames = len(time_series)

def generate_frame(i, frame_time):
    """生成单帧图像并保存"""
    # 更新所有对象位置
    constellation.update_all(frame_time)
    sats_gps = constellation.get_all_current_gps()
    sats_lon = [g[0] for g in sats_gps]
    sats_lat = [g[1] for g in sats_gps]
    
    uav_1.update(frame_time)
    uav_2.update(frame_time)
    uavs_lon = [uav_1.current_gps[0], uav_2.current_gps[0]]
    uavs_lat = [uav_1.current_gps[1], uav_2.current_gps[1]]
    
    # 创建地图
    import warnings
    warnings.filterwarnings("ignore", message="facecolor will have no effect as it has been defined as \"never\".")
    
    fig = plt.figure(figsize=(6, 4))
    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
    ax.set_extent([-180, 180, -90, 90], crs=ccrs.PlateCarree())
    
    # 添加地图特征
    ax.add_feature(cfeature.LAND, facecolor='grey', edgecolor='none', alpha=0.7)
    ax.add_feature(cfeature.OCEAN, facecolor='#A6CAE0', edgecolor='none')
    ax.add_feature(cfeature.COASTLINE, linewidth=0.1, color='white')
    
    # 绘制网格线
    gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                      linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
    gl.top_labels = False
    gl.right_labels = False
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    
    # 绘制点
    ax.scatter(sats_lon, sats_lat, s=20, alpha=0.7, c="green", edgecolors="black", label="Satellites", transform=ccrs.PlateCarree())
    ax.scatter(stations_lon_list, stations_lat_list, s=20, alpha=0.7, c="blue", edgecolors="black", label="Ground Stations", transform=ccrs.PlateCarree())
    ax.scatter(uavs_lon, uavs_lat, s=20, alpha=0.7, c="red", edgecolors="black", label="UAVs", transform=ccrs.PlateCarree())
    
    # 绘制链路
    for dev_send in ALL_DEVICES:
        base_send = dev_send.platform
        radio_recvs = dev_send.get_current_access_radios()
        for radio_recv in radio_recvs:
            base_recv = radio_recv.platform
            try:
                ax.plot(
                    [base_send.current_gps[0], base_recv.current_gps[0]],
                    [base_send.current_gps[1], base_recv.current_gps[1]],
                    color="black", linewidth=0.1, transform=ccrs.Geodetic()
                )
            except:
                continue
    
    # ax.legend(loc='upper right', fontsize='small')
    plt.tight_layout()
    # 保存帧（按序号命名确保顺序）
    frame_path = f"frames/frame_{i:04d}.png"
    plt.savefig(frame_path, dpi=150)
    plt.close(fig)  # 关闭图像释放内存
    return frame_path

print("INFO: 并行生成动画帧...")
# 并行生成所有帧（n_jobs=-1表示使用所有CPU核心）
with Parallel(n_jobs=-1, verbose=10) as parallel:
    # 为每个帧分配索引和时间
    frame_paths = parallel(delayed(generate_frame)(i, t) for i, t in enumerate(time_series))

# endregion

# region 拼接帧为GIF
print("INFO: 拼接帧为GIF...")
# 按顺序读取所有帧
frames = []
for path in sorted(frame_paths):  # 确保按序号排序
    frames.append(imageio.imread(path))

# 保存为GIF
imageio.mimsave('SAGIN.gif', frames, fps=10, loop=0)
print("INFO: GIF生成完成！")

# 删除临时帧文件
import shutil
shutil.rmtree("frames")

# endregion